/**
 * zephyr spi flash driver
 *   Macronix MX35XX Serial Nand Flash Memory
 *   https://www.mxic.com.tw/Lists/Datasheet/Attachments/8198/MX35LF2G14AC, 3V,
 * 2Gb, v1.1.pdf zephyr style driver
 */
#define DT_DRV_COMPAT mxicy_mx35

#include <zephyr/device.h>
#include <zephyr/drivers/spi.h>
#include <zephyr/kernel.h>
#include <zephyr/logging/log.h>
#include <zephyr/pm/device.h>
#include <zephyr/sys/byteorder.h>

#include "mx35_cmd.h"
#include "spi_flash_mx35.h"

LOG_MODULE_REGISTER(spi_flash_mx35, CONFIG_SPI_FLASH_MX35_LOG_LEVEL);

#define DEF_BUF_SET(_name, _buf_array)                                                             \
    const struct spi_buf_set _name = {                                                             \
        .buffers = _buf_array,                                                                     \
        .count = ARRAY_SIZE(_buf_array),                                                           \
    }

#define SPI_FLASH_MX35_INST_ID 0

struct flash_mx35_data {
    uint8_t state;
};

struct flash_mx35_config {
    struct spi_dt_spec bus;
    const struct gpio_dt_spec wp;
    const struct gpio_dt_spec hold;

    uint32_t chip_size;
    uint32_t page_size;
    uint32_t block_size;
    uint32_t page_oob_size;
    uint8_t jedec_id[2]; // Byte0: manufacturer ID; Byte1: Unique ID
};

static struct mx35_flash_parameters flash_parameters;

static struct flash_mx35_data *get_dev_data(const struct device *dev)
{
    return dev->data;
}

static const struct flash_mx35_config *get_dev_config(const struct device *dev)
{
    return dev->config;
}

/**************************************************************/
/* mx35 Flash commands */
static bool check_status(const struct device *dev, uint8_t check_flag);
static bool wait_until_ready(const struct device *dev);
static int set_write_enable(const struct device *dev);
static int disable_write_protect(const struct device *dev);
static int cmd_get_feature(const struct device *dev, uint8_t addr, uint8_t *data);
static int cmd_set_feature(const struct device *dev, uint8_t addr, uint8_t value);
static int cmd_read_id(const struct device *dev, uint8_t *data, uint8_t len);
static int cmd_perform_read(const struct device *dev, uint16_t flash_addr, uint8_t *data,
                            uint32_t len, uint8_t wrap_flag);
static int cmd_perform_program(const struct device *dev, uint32_t flash_addr, uint8_t *data,
                               uint32_t len, uint8_t wrap_flag);
static int cmd_block_erase(const struct device *dev, uint32_t flash_addr);

static bool check_status(const struct device *dev, uint8_t check_flag)
{
    uint8_t status = 0;
    cmd_get_feature(dev, FEATURE_STATUS, &status);
    if ((status & check_flag) == check_flag) {
        return BUSY;
    } else {
        return READY;
    }
}

static bool wait_until_ready(const struct device *dev)
{
    for (size_t i = 0; i < CONFIG_SPI_FLASH_MX35_READY_TIMEOUT; i++) {
        if (check_status(dev, SR0_OIP) == READY) {
            return READY;
        }
        k_msleep(1);
    }
    return TIMEOUT;
}

static int set_write_enable(const struct device *dev)
{
    int err;
    const struct flash_mx35_config *cfg = get_dev_config(dev);
    uint8_t const op_and_addr[] = { FLASH_CMD_WREN };
    const struct spi_buf tx_buf[] = {
        {
            .buf = (void *)&op_and_addr,
            .len = sizeof(op_and_addr),
        },
    };
    DEF_BUF_SET(tx_buf_set, tx_buf);
    err = spi_write_dt(&cfg->bus, &tx_buf_set);
    if (err) {
        LOG_ERR("SPI transaction failed with code: %d/%u", err, __LINE__);
        return FLASH_DATA_INVALID;
    }

    /* disable write protect */
    disable_write_protect(dev);

    return err;
}

static int disable_write_protect(const struct device *dev)
{
    uint8_t status = 0;
    cmd_get_feature(dev, FEATURE_BLOCK_PROTECTION, &status);
    if (status & 0x38) {
        LOG_DBG("Disable block protect bit...");
        cmd_set_feature(dev, FEATURE_BLOCK_PROTECTION, status & 0x87);
        cmd_get_feature(dev, FEATURE_BLOCK_PROTECTION, &status);
        LOG_DBG("read %02X: %02X", FEATURE_BLOCK_PROTECTION, status);
    }
    return 0;
}

static int cmd_get_feature(const struct device *dev, uint8_t addr, uint8_t *data)
{
    int err;
    const struct flash_mx35_config *cfg = get_dev_config(dev);
    uint8_t const op_and_addr[] = { FLASH_CMD_GET_FEATURE, addr };
    const struct spi_buf tx_buf[] = {
        {
            .buf = (void *)&op_and_addr,
            .len = sizeof(op_and_addr),
        },
    };
    const struct spi_buf rx_buf[] = {
        {
            .len = sizeof(op_and_addr),
        },
        {
            .buf = (void *)data,
            .len = sizeof(*data),
        },
    };
    DEF_BUF_SET(tx_buf_set, tx_buf);
    DEF_BUF_SET(rx_buf_set, rx_buf);
    err = spi_transceive_dt(&cfg->bus, &tx_buf_set, &rx_buf_set);
    if (err) {
        LOG_ERR("SPI transaction failed with code: %d/%u", err, __LINE__);
        return FLASH_DATA_INVALID;
    }
    return err;
}

static int cmd_set_feature(const struct device *dev, uint8_t addr, uint8_t value)
{
    int err;
    const struct flash_mx35_config *cfg = get_dev_config(dev);
    uint8_t const op_and_addr[] = { FLASH_CMD_SET_FEATURE, addr, value };
    const struct spi_buf tx_buf[] = {
        {
            .buf = (void *)&op_and_addr,
            .len = sizeof(op_and_addr),
        },
    };
    DEF_BUF_SET(tx_buf_set, tx_buf);
    err = spi_write_dt(&cfg->bus, &tx_buf_set);
    if (err) {
        LOG_ERR("SPI transaction failed with code: %d/%u", err, __LINE__);
        return FLASH_DATA_INVALID;
    }
    return err;
}

static int cmd_read_id(const struct device *dev, uint8_t *data, uint8_t len)
{
    if (check_status(dev, SR0_OIP) != READY) {
        return FLASH_BUSY;
    }

    int err;
    const struct flash_mx35_config *cfg = get_dev_config(dev);

    uint8_t const op_and_addr[] = { FLASH_CMD_RDID, 0x00 };
    const struct spi_buf tx_buf[] = {
        {
            .buf = (void *)&op_and_addr,
            .len = sizeof(op_and_addr),
        },
    };
    const struct spi_buf rx_buf[] = {
        {
            .len = sizeof(op_and_addr),
        },
        {
            .buf = (void *)data,
            .len = len,
        },
    };
    DEF_BUF_SET(tx_buf_set, tx_buf);
    DEF_BUF_SET(rx_buf_set, rx_buf);
    err = spi_transceive_dt(&cfg->bus, &tx_buf_set, &rx_buf_set);
    if (err) {
        LOG_ERR("SPI transaction failed with code: %d/%u", err, __LINE__);
        return FLASH_DATA_INVALID;
    }
    return err;
}

static int cmd_perform_read(const struct device *dev, uint16_t flash_addr, uint8_t *data,
                            uint32_t len, uint8_t wrap_flag)
{
    if (flash_addr > FLASH_SIZE) {
        return FLASH_ADDR_INVALID;
    }

    if (check_status(dev, SR0_OIP) != READY) {
        LOG_WRN("flash is busy");
        return FLASH_BUSY;
    }

    int err;
    const struct flash_mx35_config *cfg = get_dev_config(dev);

    /* Read data from flash to page cache */
    {
        uint32_t row_addr = flash_addr >> 12;
        uint8_t const op_and_addr[] = {
            FLASH_CMD_READ,
            (row_addr >> 16) & 0xFF,
            (row_addr >> 8) & 0xFF,
            row_addr & 0xFF,
        };
        const struct spi_buf tx_buf[] = {
            {
                .buf = (void *)&op_and_addr,
                .len = sizeof(op_and_addr),
            },
        };
        DEF_BUF_SET(tx_buf_set, tx_buf);
        err = spi_write_dt(&cfg->bus, &tx_buf_set);
        if (err) {
            LOG_ERR("SPI transaction failed with code: %d/%u", err, __LINE__);
            return FLASH_DATA_INVALID;
        }

        if (wait_until_ready(dev) != READY) {
            return FLASH_OPERATION_TIMEOUT;
        }
    }

    /* Read data from page cache to user buffer */
    {
        uint16_t col_addr = flash_addr & 0x0FFF;
        uint8_t const op_and_addr[] = {
            FLASH_CMD_READ_CACHE,
            (wrap_flag << 4) | ((col_addr >> 8) & 0x0F),
            col_addr & 0xFF,
            FLASH_CMD_DUMMY_DATA,
        };
        const struct spi_buf tx_buf[] = {
            {
                .buf = (void *)&op_and_addr,
                .len = sizeof(op_and_addr),
            },
        };
        const struct spi_buf rx_buf[] = {
            {
                .len = sizeof(op_and_addr),
            },
            {
                .buf = (void *)data,
                .len = len,
            },
        };
        DEF_BUF_SET(tx_buf_set, tx_buf);
        DEF_BUF_SET(rx_buf_set, rx_buf);
        err = spi_transceive_dt(&cfg->bus, &tx_buf_set, &rx_buf_set);
        if (err) {
            LOG_ERR("SPI transaction failed with code: %d/%u", err, __LINE__);
            return FLASH_DATA_INVALID;
        }
    }

    return err;
}

static int cmd_perform_program(const struct device *dev, uint32_t flash_addr, uint8_t *data,
                               uint32_t len, uint8_t wrap_flag)
{
    if (check_status(dev, SR0_OIP) != READY) {
        LOG_WRN("flash is busy");
        return FLASH_BUSY;
    }

    set_write_enable(dev);

    int err;
    const struct flash_mx35_config *cfg = get_dev_config(dev);

    /* copy data to cache from user buffer */
    {
        uint16_t col_addr = flash_addr & 0x0FFF;
        uint8_t const op_and_addr[] = {
            FLASH_CMD_PP_LOAD,
            (wrap_flag << 4) | ((col_addr >> 8) & 0x0F),
            col_addr & 0xFF,
        };
        const struct spi_buf tx_buf[] = {
            {
                .buf = (void *)&op_and_addr,
                .len = sizeof(op_and_addr),
            },
            {
                .buf = (void *)data,
                .len = len,
            },
        };
        DEF_BUF_SET(tx_buf_set, tx_buf);
        err = spi_write_dt(&cfg->bus, &tx_buf_set);
        if (err) {
            LOG_ERR("SPI transaction failed with code: %d/%u", err, __LINE__);
            return FLASH_DATA_INVALID;
        }
    }

    /* commit data to flash from cache */
    {
        uint32_t row_addr = flash_addr >> 12;
        uint8_t const op_and_addr[] = {
            FLASH_CMD_PROGRAM_EXEC,
            (row_addr >> 16) & 0xFF,
            (row_addr >> 8) & 0xFF,
            row_addr & 0xFF,
        };
        const struct spi_buf tx_buf[] = {
            {
                .buf = (void *)&op_and_addr,
                .len = sizeof(op_and_addr),
            },
        };
        DEF_BUF_SET(tx_buf_set, tx_buf);
        err = spi_write_dt(&cfg->bus, &tx_buf_set);
        if (err) {
            LOG_ERR("SPI transaction failed with code: %d/%u", err, __LINE__);
            return FLASH_DATA_INVALID;
        }

        if (wait_until_ready(dev) == READY) {
            uint8_t status = 0;
            cmd_get_feature(dev, FEATURE_STATUS, &status);
            if ((status & SR2_PROGRAM_FAILED) == SR2_PROGRAM_FAILED) {
                return FLASH_PROGRAM_FAILED;
            }
        } else {
            return FLASH_OPERATION_TIMEOUT;
        }
    }

    return err;
}

static int cmd_block_erase(const struct device *dev, uint32_t flash_addr)
{
    if (flash_addr > FLASH_SIZE) {
        return FLASH_ADDR_INVALID;
    }

    if (check_status(dev, SR0_OIP) != READY) {
        LOG_WRN("flash is busy");
        return FLASH_BUSY;
    }

    set_write_enable(dev);

    int err;
    const struct flash_mx35_config *cfg = get_dev_config(dev);
    uint32_t row_addr = flash_addr >> 12;
    uint8_t const op_and_addr[] = {
        FLASH_CMD_BE,
        (row_addr >> 16) & 0xFF,
        (row_addr >> 8) & 0xFF,
        row_addr & 0xFF,
    };
    const struct spi_buf tx_buf[] = {
        {
            .buf = (void *)&op_and_addr,
            .len = sizeof(op_and_addr),
        },
    };
    DEF_BUF_SET(tx_buf_set, tx_buf);
    err = spi_write_dt(&cfg->bus, &tx_buf_set);
    if (err) {
        LOG_ERR("SPI transaction failed with code: %d/%u", err, __LINE__);
        return FLASH_DATA_INVALID;
    }

    if (wait_until_ready(dev) == READY) {
        uint8_t status = 0;
        cmd_get_feature(dev, FEATURE_STATUS, &status);
        if ((status & SR2_ERASE_FAILED) == SR2_ERASE_FAILED) {
            return FLASH_ERASE_FAILED;
        } else {
            return FLASH_SUCCESS;
        }
    } else {
        return FLASH_OPERATION_TIMEOUT;
    }
}

/**
 * @brief translate logic addr to flash write addr
 *  - RA[16:6], Block address
 *  - RA[5:0], Page address
 *  - CA[11:0], Byte address
 * @param addr
 * @return
 */
static uint32_t flash_addr_convert(uint32_t addr)
{
    uint16_t col_addr = 0;
    uint16_t page_nbr = 0;
    uint32_t block_nbr = 0;
    uint32_t flash_addr = 0;

    /* Column Addr, CA[11:0] */
    col_addr = addr % 2048;
    /* numbers of page */
    page_nbr = addr / 2048;
    /* numbers of block */
    block_nbr = page_nbr / 64;

    flash_addr = block_nbr << 18;
    flash_addr |= (page_nbr % 64) << 12;
    flash_addr |= col_addr;

    return flash_addr;
}
/**************************************************************/

static int check_jedec_id(const struct device *dev)
{
    const struct flash_mx35_config *cfg = get_dev_config(dev);
    const uint8_t *expected_id = cfg->jedec_id;
    uint8_t read_id[sizeof(cfg->jedec_id)];

    cmd_read_id(dev, read_id, sizeof(read_id));
    if (memcmp(expected_id, read_id, sizeof(read_id)) != 0) {
        LOG_ERR("Wrong JEDEC ID: %02X %02X, "
                "expected: %02X %02X",
                read_id[0], read_id[1], expected_id[0], expected_id[1]);
        return -ENODEV;
    }

    LOG_DBG("Flash JEDEC ID: %02X %02X", expected_id[0], expected_id[1]);
    return 0;
}

static int flash_mx35_read(const struct device *dev, uint32_t addr, uint8_t *data, uint32_t len,
                           uint8_t wrap_flag)
{
    int err;
    uint32_t flash_addr = 0;
    uint32_t read_addr = addr;
    size_t btr_avl = 0;
    size_t btr_len = 0;
    size_t btr_remain = len;
    size_t rd_pos = 0;
    uint8_t *p_data = (uint8_t *)data;
    const struct flash_mx35_config *cfg = get_dev_config(dev);

    LOG_DBG("read addr %d in flash block %d, page %d, offset %d, plane %d", read_addr,
            (read_addr / cfg->block_size), (read_addr / cfg->page_size),
            (read_addr % cfg->page_size), (read_addr / cfg->block_size) & 0x1);

    // 计算当前页可读大小
    btr_avl = cfg->page_size - ((read_addr % cfg->page_size) & 0x0FFF);
    do {
        if (btr_remain > btr_avl) {
            btr_len = btr_avl;
        } else {
            btr_len = btr_remain;
        }

        // 读操作
        flash_addr = flash_addr_convert(read_addr);
        err = cmd_perform_read(dev, flash_addr, (p_data + rd_pos), btr_len,
                               ((flash_addr >> 18) & 0x1));
        if (err) {
            LOG_ERR("flash read failed with err code: %d/%u.", err, __LINE__);
            return err;
        }

        btr_remain -= btr_len;
        read_addr += btr_len;
        rd_pos += btr_len;
    } while (btr_remain > 0);

    return err;
}

static int flash_mx35_write(const struct device *dev, uint32_t addr, const uint8_t *data,
                            uint32_t len, uint8_t wrap_flag)
{
    int err = 0;
    uint32_t flash_addr = 0;
    uint32_t btw_len = len;
    uint32_t write_addr = addr;
    int btw_remain = len;
    uint32_t w_pos = 0;
    uint8_t *p_data = (uint8_t *)data;

    const struct flash_mx35_config *cfg = get_dev_config(dev);

    LOG_DBG("write addr %d in flash block %d, page %d, offset %d, plane %d", write_addr,
            (write_addr / cfg->block_size), (write_addr / cfg->page_size),
            (write_addr % cfg->page_size), (write_addr / cfg->block_size) & 0x1);

    do {
        // erase first if in a new block address
        if ((write_addr % cfg->block_size) == 0) {
            uint32_t block_num = 0;
            uint32_t erase_flash_addr = 0;
            block_num = write_addr / cfg->block_size;
            erase_flash_addr = block_num << 18;

            LOG_DBG("addr %d, block id %d in plane %d", write_addr, block_num, block_num & 0x01);
            LOG_DBG("++++++FLASH ERASE++++++");

            uint8_t status = 0;
            cmd_get_feature(dev, FEATURE_BLOCK_PROTECTION, &status);
            LOG_DBG("Get Feature: A0 -> %02X", status);
            if (status & 0x38) {
                LOG_DBG("Disable block protect bit...");
                cmd_set_feature(dev, FEATURE_BLOCK_PROTECTION, status & 0x87);
            }

            LOG_DBG("Block Erasing block addr %d", erase_flash_addr);
            cmd_block_erase(dev, erase_flash_addr);
        }

        uint32_t av = cfg->page_size - (write_addr % cfg->page_size);
        if (btw_remain > av) {
            btw_len = av;
        } else {
            btw_len = btw_remain;
        }

        flash_addr = flash_addr_convert(write_addr);
        err = cmd_perform_program(dev, flash_addr, (p_data + w_pos), btw_len,
                                  (flash_addr >> 18) & 0x1);
        if (err) {
            LOG_ERR("flash write failed with err code: %d/%u", err, __LINE__);
            return err;
        }

        btw_remain -= btw_len;
        write_addr += btw_len;
        w_pos += btw_len;

        if (write_addr >= cfg->chip_size) {
            write_addr = 0;
        }

    } while (btw_remain > 0);

    return (err != 0) ? -EIO : 0;
}

static int flash_mx35_erase(const struct device *dev, uint32_t addr)
{
    LOG_DBG("block erase addr %d", addr);
    return cmd_block_erase(dev, addr);
}

static int flash_mx35_get_state(const struct device *dev)
{
    struct flash_mx35_data *p_data = get_dev_data(dev);
    return p_data->state;
}

static const struct mx35_flash_parameters *flash_mx35_get_parameters(const struct device *dev)
{
    const struct flash_mx35_config *p_cfg = get_dev_config(dev);
    flash_parameters.chip_size = p_cfg->chip_size;
    flash_parameters.block_size = p_cfg->block_size;
    flash_parameters.page_size = p_cfg->page_size;
    flash_parameters.page_oob_size = p_cfg->page_oob_size;
    memcpy(flash_parameters.jedec_id, p_cfg->jedec_id, sizeof(p_cfg->jedec_id));
    return &flash_parameters;
}

static int flash_mx35_init(const struct device *dev)
{
    int err = 0;
    const struct flash_mx35_config *p_cfg = get_dev_config(dev);
    struct flash_mx35_data *p_data = get_dev_data(dev);

    p_data->state = 0;

    if (!spi_is_ready_dt(&p_cfg->bus)) {
        LOG_ERR("SPI bus %s not ready.", p_cfg->bus.bus->name);
        return -ENODEV;
    }

    if (device_is_ready(p_cfg->wp.port)) {
        if (gpio_pin_configure_dt(&p_cfg->wp, GPIO_OUTPUT_INACTIVE)) {
            LOG_ERR("Couldn't configure write protect pin");
            return -ENODEV;
        }
        gpio_pin_set(p_cfg->wp.port, p_cfg->wp.pin, 1);
    }

    if (device_is_ready(p_cfg->hold.port)) {
        if (gpio_pin_configure_dt(&p_cfg->hold, GPIO_OUTPUT_INACTIVE)) {
            LOG_ERR("Couldn't configure hold pin");
            return -ENODEV;
        }
        gpio_pin_set(p_cfg->hold.port, p_cfg->hold.pin, 1);
    }

    err = check_jedec_id(dev);
    if (err) {
        LOG_ERR("flash MX35 driver init FAILED.\n");
        return err;
    }

    p_data->state = 1;

    LOG_INF("flash MX35 driver init OK.\n");
    return err;
}

static struct flash_mx35_data _data = {
    .state = 0,
};

static const struct flash_mx35_config _config = {
    .bus = {
        .bus = DEVICE_DT_GET(DT_BUS(DT_DRV_INST(SPI_FLASH_MX35_INST_ID))),
        .config = {
            .operation = SPI_OP_MODE_MASTER | SPI_TRANSFER_MSB | SPI_WORD_SET(8) | SPI_LINES_SINGLE,
            .frequency = DT_INST_PROP(SPI_FLASH_MX35_INST_ID, spi_max_frequency),
            .cs = SPI_CS_CONTROL_INIT(DT_NODELABEL(mx35lf2g14ac), 0),
        },
    },
    .wp = GPIO_DT_SPEC_GET(DT_DRV_INST(SPI_FLASH_MX35_INST_ID), wp_gpios),
    .hold = GPIO_DT_SPEC_GET(DT_DRV_INST(SPI_FLASH_MX35_INST_ID), hold_gpios),
    .chip_size = DT_INST_PROP(SPI_FLASH_MX35_INST_ID, size),                                                      
    .page_size = DT_INST_PROP(SPI_FLASH_MX35_INST_ID, page_size),                                                 
    .page_oob_size = DT_INST_PROP(SPI_FLASH_MX35_INST_ID, page_oob_size),                                         
    .block_size = DT_INST_PROP(SPI_FLASH_MX35_INST_ID, block_size),                                               
    .jedec_id = DT_INST_PROP(SPI_FLASH_MX35_INST_ID, jedec_id), 
};

static const struct mx35_flash_driver_api _apis = {
    .read = flash_mx35_read,
    .write = flash_mx35_write,
    .erase = flash_mx35_erase,
    .get_state = flash_mx35_get_state,
    .get_parameters = flash_mx35_get_parameters,
};

DEVICE_DT_DEFINE(DT_NODELABEL(mx35lf2g14ac), &flash_mx35_init, NULL, &_data, &_config, POST_KERNEL,
                 CONFIG_SPI_FLASH_MX35_INIT_PRIORITY, &_apis);
